from stock import Stock

class Strategy(Stock):
    def __init__(self, data):
        Stock.__init__(self, data)
    """
    单因子策略：涨幅策略
    选出回测时间段内涨幅最大的前num支股票
    """
    def choose(self, date, num = 2, period = 3):
        buy_list = []
        growth = {}
        # 获得period天前的日期
        before_date = self.dates[self.dates.index(date) - period]
        for id in self.ids:
            # 收盘价
            close_price = self.close_price(id, date)
            close_price_before = self.close_price(id, before_date)
            # 收盘价的涨幅
            growth[id] = (close_price - close_price_before)/close_price_before
        # 对收盘价涨幅进行排序
        growth_order = sorted(growth.items(),key=lambda x:x[1],reverse=True)
        # 返回上涨幅度前num大的股票代码
        for i in range(num):
            # 只选上涨的股票
            if growth_order[i][1] > 0:
                buy_list.append(growth_order[i][0])
        return buy_list
    """
    多因子策略：RSI策略与AR（人气指标）策略结合
    RSI相对强弱指标是一个信号指标。
    AR人气指标通过一定时期内开盘价、最高价、以及最低价之间的关系，来分析多空力量的对比，反映市场买卖人气，分析价格波动，达到追踪价格未来动向的目的。
    """
    def choose2(self, date, period = 20):
        buy_list = []
        AR = {}
        for id in self.ids:
            high_open = 0
            open_low = 0
            for d in range(1, period + 1):
                date = self.dates[self.dates.index(date) - d]
                # 获得开盘价、最高价和最低价
                open_price = self.open_price(id, date)
                high_price = self.high_price(id, date)
                low_price = self.low_price(id, date)
                high_open += high_price - open_price
                open_low += open_price - low_price
            # 计算公式为 AR = [N天所有（High-Open）的和/ N天所有（Open—Low）的和] * 100
            AR[id] = high_open  * 100 / open_low
            # AR低于超卖线80，RSI低于40时则买入
            if AR[id] < 80 and self.rsi(id, date) < 40:
                buy_list.append(id)
        return buy_list

